from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
import uvicorn
from PyCmpltrtok.common import md5, uuid
import redis
import time
import asyncio
import json

app = FastAPI()
rdb = redis.Redis('127.0.0.1', 6379, 0)
PWD = 'my_password_001'

THRESH = 100


@app.post("/api")
async def do_infer(request: Request):
    # 接收输入
    req_json = await request.json()  # 请求json

    async def get_generator():

        xinput = req_json['input']
        xusername = req_json['username']

        # 防止蹭模型
        xcheck = req_json['check']
        xmy_check = md5(xinput + PWD + xusername)
        if xcheck.lower() != xmy_check.lower():
            print('Checking not passed!')
            return

        print('input:', xinput)

        # 数据放入队列和hash
        xuuid = uuid()
        rdb.hset('uuid2input', xuuid, xinput.encode('utf8'))
        rdb.hset('uuid2username', xuuid, xusername.encode('utf8'))
        rdb.lpush('queue', xuuid)

        # 等着daemon处理

        # 轮询hash以拿回结果
        prev = None
        repeat_count = 0
        while True:
            xoutput = rdb.get('chat_cache_' + xusername)
            if xoutput is None:
                await asyncio.sleep(0.001)  # 重要
                continue
            final_output = rdb.hget('uuid2output', xuuid)
            if final_output is not None:
                print('OVER')
                final_output = final_output.decode('utf8')
                break
            xoutput = xoutput.decode('utf8')
            if prev is not None and prev == xoutput:
                repeat_count += 1
                print(f'THE SAME x{repeat_count}')
                if repeat_count > THRESH:
                    print('THRESH!')
                    return
                await asyncio.sleep(0.001)  # 重要
                continue
            else:
                prev = xoutput
                repeat_count = 0
            print('output:', xoutput)
            # 返回输出
            res_dict = dict()
            res_dict['input'] = xinput
            res_dict['output'] = xoutput
            yield json.dumps(res_dict).encode('utf8') + b"\0"

        # 清理
        rdb.hdel('uuid2input', xuuid)
        rdb.hdel('uuid2output', xuuid)
        # 返回输出
        res_dict = dict()
        res_dict['input'] = xinput
        res_dict['output'] = final_output
        yield json.dumps(res_dict).encode('utf8') + b"\0"
        
    generator = get_generator()
    return StreamingResponse(generator)


if '__main__' == __name__:
    # https://stackoverflow.com/questions/63177681/is-there-a-difference-between-running-fastapi-from-uvicorn-command-in-dockerfile
    uvicorn.run('fastapi_service:app', host='0.0.0.0', port=7766)
